package com.jboltai.capability.vdb.types.postgresql;

import cn.hutool.core.util.StrUtil;
import com.jboltai.resource.embedding.EmbeddingModel;
import com.jboltai.resource.vdb.VDBResource;
import com.jboltai.resource.vdb.VDBType;
import com.pgvector.PGvector;
import com.zaxxer.hikari.HikariDataSource;

import java.sql.Connection;

public class PostgreSQLVDBResource extends VDBResource {
    private static HikariDataSource ds;

    public PostgreSQLVDBResource(String id, String url, String username, String password, EmbeddingModel embeddingModel) {
        super(id, VDBType.PGVECTOR, url, username, password, embeddingModel);
    }

    public Connection getConnection() {
        try {
            initDataSource(url, username, password);
            Connection conn = ds.getConnection();
            PGvector.addVectorType(conn);
            return conn;
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    private void initDataSource(String url, String username, String password) {
        if (ds == null) {
            synchronized (PostgreSQLVDBResource.class) {
                if (ds == null) {
                    String database = getDatabase();
                    if (StrUtil.isBlank(database)) {
                        return;
                    }
                    ds = new HikariDataSource();
                    ds.setJdbcUrl(url + database);
                    ds.setUsername(username);
                    ds.setPassword(password);
                }
            }
        }
    }
}
